#include <GrabcutFilter.h>
#include <Picture.h>

#include <highgui.h>
#include <cv.h>
#include <iostream>
#include <vector>

using namespace std;

void getBinMask( const cv::Mat& comMask, cv::Mat& binMask )
{
    if( comMask.empty() || comMask.type()!=CV_8UC1 )
        CV_Error( CV_StsBadArg, "comMask is empty or has incorrect type (not CV_8UC1)" );
    if( binMask.empty() || binMask.rows!=comMask.rows || binMask.cols!=comMask.cols )
        binMask.create( comMask.size(), CV_8UC1 );
    binMask = comMask & 1;
}


bool GrabcutFilter::init() {
    if (!picture.isInitialized()) return false;
    
    initMask();
    return true;    
}

void GrabcutFilter::initMask () {
    /* step 1 - create mask */
    mask.create(picture.getMatrix().size(),CV_8UC1);

    /* step 2 - initialize matrix as background */
    mask.setTo(cv::GC_BGD);
 
    /* step 3 - iterate over the rectangles and set the as foreground */
    for(vector<cv::Rect>::const_iterator r = ROI.begin(); r != ROI.end(); r++) {
        (mask(*r)).setTo(cv::Scalar(cv::GC_PR_FGD));
    }
}

GrabcutFilter::GrabcutFilter(const Picture& pic, vector<cv::Rect>& ROI):
    picture(pic),
    ROI(ROI),
    isValidated(false),
    iterCount (0),
    state (GrabcutFilter::NOT_SET)
{
    initialized = init();
}

GrabcutFilter::~GrabcutFilter() {
    destroy();
}
void GrabcutFilter::destroy() {
}

bool GrabcutFilter::refresh(cv::Mat& image) const {
    cv::Mat bin_mask(picture.getMatrix().size(), CV_8UC1, cv::Scalar(0));
    cv::Mat bin_mask_inv;
    cv::Mat gray;
    cv::Mat gray_part;
    cv::Mat color_part;

    // gray version of the entire image (one channel);
    cv::cvtColor(picture.getMatrix(), gray, CV_BGR2GRAY);
    if (!isValidated) {
        // mark selected ROI in mask
        for(vector<cv::Rect>::const_iterator r = ROI.begin(); r != ROI.end(); r++) {
            cv::Mat roi(bin_mask, *r);
            roi |= cv::Scalar(1);
        }
    } else {
        getBinMask(mask, bin_mask);
    }

    // invert zero's with ones and vice versa;
    bin_mask_inv = bin_mask ^ cv::Scalar(1);

    picture.getMatrix().copyTo(color_part, bin_mask);
    gray.copyTo(gray_part, bin_mask_inv);

    vector<cv::Mat> mats;
    // split channels to Red Green and Blue;
    cv::split(color_part, mats);
    // add gray part to each channel
    mats[0] += gray_part;
    mats[1] += gray_part;
    mats[2] += gray_part;
    // sett'oo
    cv::merge(mats, image);

    return true;
}

int GrabcutFilter::nextIter(void* moreData) {
    cv::Rect dummy;
    cv::Mat original = picture.getMatrix();
    if(moreData){
        cv::Mat userInteraction = *(cv::Mat*)moreData;
        for(int y=0;y<mask.rows;y++)
        {
            for(int x=0;x<mask.cols;x++)
            {
                if( userInteraction.at<uchar>(y,x) != NOTHING )
                    mask.at<uchar>(y,x) = userInteraction.at<uchar>(y,x);
            }
        }
    }
    grabCut(original,
            mask,
            dummy,
            bgdModel, fgdModel,
            1, cv::GC_INIT_WITH_MASK);
    
    // set validation flag until user makes his next refine.    
    isValidated = true;

    return iterCount;
}


vector< vector < cv::Mat > > GrabcutFilter::exportROIs()
{
    cv::Mat originalMask, originalPicture;
    getBinMask(mask, originalMask);
    picture.getMatrix().copyTo(originalPicture);
    vector < vector<cv::Mat> > output;
    for(vector<cv::Rect>::const_iterator r = ROI.begin(); r != ROI.end(); r++) {
        cv::Mat resMask(originalMask, *r);
        cv::Mat resPicture(originalPicture, *r);
        vector < cv::Mat > currentResult;
        currentResult.push_back(resMask);
        currentResult.push_back(resPicture);
        output.push_back(currentResult);
    }   
    return output;
}

boost::shared_ptr<Picture> GrabcutFilter::exportPicture() {
    cv::Mat bin_mask;
    getBinMask(mask, bin_mask);

    cv::Mat res; // notice that this object will "die" as the function returns,
                  // but Picture's c'tor will make his own copy of the object.
    picture.getMatrix().copyTo(res, bin_mask);

    return boost::shared_ptr<Picture>(new Picture(res));
}
